import json
import logging
from typing import Dict, List
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from starlette.websockets import WebSocketState
import asyncio
from sqlalchemy.orm import Session, joinedload
from fastapi import Depends, HTTPException, APIRouter

from backend.app import models
from backend.app.database import SessionLocal

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()
class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[int, List[WebSocket]] = {}



    async def connect(self, websocket: WebSocket, vote_id: int, db: Session):
        await websocket.accept()
        if vote_id not in self.active_connections:
            self.active_connections[vote_id] = []
        self.active_connections[vote_id].append(websocket)

        # 查询数据库中的投票信息
        db_vote = db.query(models.Vote).filter(models.Vote.id == vote_id).first()
        if db_vote is None:
            raise HTTPException(status_code=404, detail="Vote not found")
        # 加载关联的选项
        options = db.query(models.VoteOption).filter(models.VoteOption.vote_id == vote_id).all()
        db_vote.options = options

        ds = {
            "question":db_vote.question,
            "id":db_vote.id,
            "question_type":db_vote.question_type,
            "options":[{"id":item.id,"content":item.content,"vote_count":item.vote_count,"vote_id":item.vote_id} for item in db_vote.options],
        }
        print(ds)
        # 广播初始化消息
        await self.broadcast(vote_id,json.dumps(ds))

    def disconnect(self, vote_id: int, websocket: WebSocket):
        if vote_id in self.active_connections:
            connections = self.active_connections[vote_id]
            if websocket in connections:
                connections.remove(websocket)
                logger.info(f"Connection removed from vote_id {vote_id}")

    async def broadcast(self, vote_id: int, message: str):
        if vote_id in self.active_connections:
            logger.info(f"Broadcasting message to vote_id {vote_id}: {message}")
            for connection in self.active_connections[vote_id][:]:
                try:
                    if connection.application_state != WebSocketState.CONNECTED:
                        logger.warning(f"Connection is not connected: {connection.client}")
                        self.disconnect(vote_id, connection)
                        continue

                    logger.info(f"Sending message to {connection.client}: {message}")
                    await connection.send_text(message)
                except Exception as e:
                    logger.error(f"Failed to send message to {connection.client}: {e}")
                    self.disconnect(vote_id, connection)

router = APIRouter()
cm = ConnectionManager()
@router.websocket("/ws/votes/{vote_id}")
async def vote_websocket(websocket: WebSocket, vote_id: int,db: Session = Depends(get_db)):
    try:
        await cm.connect(websocket, vote_id,db)
        async for data in websocket.iter_text():
            await process_vote(vote_id, "{name:你好}")
    except WebSocketDisconnect:
        logger.info("WebSocket connection closed")
        cm.disconnect(vote_id, websocket)
    except Exception as e:
        logger.error(f"Unexpected error in WebSocket handler: {e}")
        cm.disconnect(vote_id, websocket)
    finally:
        if websocket.application_state == WebSocketState.CONNECTING:
            await websocket.close()

async def process_vote(vote_id: int, vote_data: str):
    try:
        logger.info(f"Processing vote data: {vote_data}")
        # 实现投票处理逻辑
        # 例如：更新数据库中的投票结果
        # 例如：广播最新的投票结果（这里使用接收到的vote_data作为示例）
        await cm.broadcast(vote_id, vote_data)
    except Exception as e:
        logger.error(f"Error processing vote: {e}")